# file: data/factor_exploder.py

import pandas as pd
from data.stock_data_loader import StockDataLoader
from datetime import timedelta


class ExploderFactorGenerator:
    def __init__(self, data_loader=None):
        self.loader = data_loader or StockDataLoader()

    def generate_zscore_single_max_money(self, start_date=None, end_date=None):
        hits = self.loader.get_hit_board_data(start_date, end_date)
        if hits.empty:
            return pd.DataFrame()

        hits["add_date"] = pd.to_datetime(hits["add_date"])
        max_money_df = hits.groupby(["stock_code", "add_date"])["money"].max().reset_index()
        max_money_df.rename(columns={"money": "single_max_money"}, inplace=True)

        max_money_df["prev_date"] = max_money_df["add_date"] - timedelta(days=1)
        prev_df = max_money_df[["stock_code", "add_date", "single_max_money"]].copy()
        prev_df.rename(columns={
            "add_date": "prev_date",
            "single_max_money": "prev_single_max_money"
        }, inplace=True)

        merged = pd.merge(max_money_df, prev_df, how="left", on=["stock_code", "prev_date"])
        merged["zscore_single_max_money"] = (
            (merged["single_max_money"] - merged["prev_single_max_money"]) /
            (merged["prev_single_max_money"].rolling(5).std())
        ).fillna(0)

        return merged

    def generate_zscore_continue_sum_money(self, start_date=None, end_date=None):
        hits = self.loader.get_hit_board_data(start_date, end_date)
        if hits.empty:
            return pd.DataFrame()

        hits["add_date"] = pd.to_datetime(hits["add_date"])
        hits = hits.sort_values(["stock_code", "batch_id"])

        result_rows = []
        for stock_code, group in hits.groupby("stock_code"):
            group = group.sort_values("batch_id")
            group["batch_time"] = pd.to_datetime(group["batch_id"].astype(str), format="%Y%m%d%H%M%S")
            group["time_diff"] = group["batch_time"].diff().dt.total_seconds().fillna(999)
            group["new_group"] = (group["time_diff"] > 60).cumsum()
            grouped_money = group.groupby(["add_date", "new_group"])["money"].sum().reset_index()
            grouped_money["stock_code"] = stock_code
            grouped_max = grouped_money.groupby(["stock_code", "add_date"])["money"].max().reset_index()
            grouped_max.rename(columns={"money": "continue_sum_money"}, inplace=True)
            result_rows.append(grouped_max)

        continue_df = pd.concat(result_rows, ignore_index=True)
        continue_df["prev_date"] = continue_df["add_date"] - timedelta(days=1)

        prev_df = continue_df[["stock_code", "add_date", "continue_sum_money"]].copy()
        prev_df.rename(columns={
            "add_date": "prev_date",
            "continue_sum_money": "prev_continue_sum_money"
        }, inplace=True)

        merged = pd.merge(continue_df, prev_df, how="left", on=["stock_code", "prev_date"])
        merged["zscore_continue_sum_money"] = (
            (merged["continue_sum_money"] - merged["prev_continue_sum_money"]) /
            (merged["prev_continue_sum_money"].rolling(5).std())
        ).fillna(0)

        return merged
